module mDiffEq

        use mPrecision
        use mObject

        implicit none

        integer, parameter :: MAX_STEPS = 10000, &    ! Maximum number of Burlisch-Stoer steps.
                              SEQUENCE(MAX_ROWS + 1) = (/ 2, 4, 6, 8, 10, 12, 14, 16 /) ! Deuflhard's sequence of successive subdivisions.

        real(D), parameter :: SAFETY_FACTOR1 = 4.0_D, &   ! Factor by which error estimates will be scaled.
                              SAFETY_FACTOR2 = 0.7_D, &   ! Extra stepsize reduction factor used in cases when we are not converging before k_max.
                              REDUCTION_MAX = 1.0e-5_D, & ! Maximum and
                              REDUCTION_MIN = 0.7_D, &    ! minimum allowed stepsize reduction factors.
                              SCALE_MAX = 0.1_D           ! 1 / SCALE_MAX is the maximum allowed stepsize increase factor.

        interface new
                module procedure new_diffEq
        end interface new

!        interface getUniverse
!                module procedure getUniverse_diffEq
!        end interface getUniverse
        interface params
                module procedure params_diffEq
        end interface

        private
        public :: new
!        public :: tDiffEq, new
!        public :: getUniverse
        public :: x, setX
        public :: y, y_all
        public :: params
        public :: statistics
        public :: integrate

contains

!        subroutine new_diffEq(diffEq, x, y, h, hMin, eps, u)
        subroutine new_diffEq(diffEq, x, y, h, hMin, eps, param)

                use mMisc

                type(tDiffEq), intent(out) :: diffEq
                real(D), intent(in) :: x    ! "Starting" x value
                real(D), intent(in) :: y(:) ! Boundary conditions at x
                real(D), intent(in), optional :: h, hMin, eps, param(:)
!                type(tUniverse), intent(in), optional :: u

                real(D), parameter :: DIFFEQ_DEFAULT_EPS = epsilon(1.0_D), &
                                      DIFFEQ_DEFAULT_H = 0.3_D, &
                                      DEFAULT_H_MIN = 0.0_D

                integer :: k

!                if (D == 8) then
!                        write(0, *) 'WARNING: diffEq not tested with double precision.'
!                        write(0, *) '         You may have to decrease eps.'
!                end if

                diffEq%x = x

                call assert( size(y) <= MAX_EQUATIONS, 'Too many equations in system. Increase MAX_EQUATIONS.' )
                diffEq%nEq = size(y)
                diffEq%y(1:diffEq%nEq) = y

                if( present(h) ) then
                        diffEq%h = h
                else
                        diffEq%h = DIFFEQ_DEFAULT_H
                end if

                if( present(hMin) ) then
                        diffEq%hMin = hMin
                else
                        diffEq%hMin = DEFAULT_H_MIN
                end if

                if( present(eps) ) then
                        diffEq%eps = eps
                else
                        diffEq%eps = DIFFEQ_DEFAULT_EPS
                end if

!                if( present(u) ) diffEq%u = u
                if( present(param) ) then
                        call assert( size(param) <= DIFFEQ_MAX_PARAMS, 'too many parameters for diffEq.' )
                        diffEq%param( 1:size(param) ) = param(:)
                end if

                diffEq%A = real( cumulativeSum(SEQUENCE, 1), D ) ! Initialize the array of work estimates.

                ! Initialize the convergence correction factors:
                where( lowerTriangular(MAX_ROWS, MAX_ROWS, -1) ) &
                        diffEq%alpha = (diffEq%eps / SAFETY_FACTOR1) ** (outerDifference( diffEq%A(2:), diffEq%A(2:) ) &
!REMOVE                                                                         / outerProduct( arithmeticVector(3.0_D, 2.0_D, MAX_ROWS), &
!                                                                                         ( diffEq%A(2:) - diffEq%A(1) + 1.0_D )))
                                                                         / outerProduct( ( diffEq%A(2:) - diffEq%A(1) + 1.0_D ), &
                                                                                           arithmeticVector(3.0_D, 2.0_D, MAX_ROWS) ))

                ! Use correction factors to determine optimal row number
                ! for convergence, as in Numerical Recipes Eq. 16.4.13.
                ! diffEq%k_max will be the row number (of the extrapolation
                ! tableau) after which the computational efficiency becomes
                ! worse:
                do k = 2, MAX_ROWS
                        if( diffEq%A(k + 1) > diffEq%A(k) * diffEq%alpha(k - 1, k) ) exit
                end do
                diffEq%k_max = k

                diffEq%k_opt = diffEq%k_max ! Set k_opt to k_max because we don't know any better yet.

                diffEq%nGood = 0
                diffEq%nBad = 0

                diffEq%first = .true.

        end subroutine new_diffEq

        pure subroutine setX(diffEq, x)
                type(tDiffEq), intent(inout) :: diffEq
                real(D), intent(in) :: x
                diffEq%x = x
        end subroutine setX

        real(D) pure function x(diffEq)
                type(tDiffEq), intent(in) :: diffEq
                x = diffEq%x
        end function x

        real(D) pure function y(diffEq)
                type(tDiffEq), intent(in) :: diffEq
                y = diffEq%y(1)
        end function y

        pure function y_all(diffEq)
                type(tDiffEq), intent(in) :: diffEq
                real(D) :: y_all(diffEq%nEq)
                y_all = diffEq%y(1:diffEq%nEq)
        end function y_all

        pure subroutine params_diffEq(diffEq, param1, param2, param3)
                type(tDiffEq), intent(in) :: diffEq
                real(D), intent(out) :: param1, param2, param3
                param1 = diffEq%param(1)
                param2 = diffEq%param(2)
                param3 = diffEq%param(3)
        end subroutine params_diffEq

        pure subroutine statistics(diffEq, nGood, nBad)
                type(tDiffEq), intent(in) :: diffEq
                integer, intent(out) :: nGood, nBad

                nGood = diffEq%nGood
                nBad = diffEq%nBad
        end subroutine statistics

        !******************************************************************
        ! Integrate from current point x stored in diffEq%x to the given x:
        !******************************************************************
        subroutine integrate(diffEq, x, derivatives)

                type(tDiffEq), intent(inout) :: diffEq
                real(D), intent(in) :: x
                interface
                        function derivatives(diffEq, x, y)
                                use mPrecision
                                use mObject
                                type(tDiffEq), intent(in) :: diffEq
                                real(D), intent(in) :: x, y(:)
                                real(D) :: derivatives( size(y) )
                        end function derivatives
                end interface

                integer :: i
                real(D) :: xStart, hDo, hDid, hNext

!                real(D), dimension( size(diffEq%y) ) :: dydx, yScal
                real(D), dimension( diffEq%nEq ) :: dydx, yScal

                xStart = diffEq%x
                hDo = sign(diffEq%h, x - xStart)

                do i = 1, MAX_STEPS

                        dydx = derivatives(diffEq, diffEq%x, diffEq%y(1:diffEq%nEq) )

                        ! This value for yScal results in constant fractional errors except very near zero crossings.
                        ! See Numerical Recipes Ch. 16.2, expecially the discussion below Eq. (16.2.8).
                        yScal(:) = abs( diffEq%y(1:diffEq%nEq) ) + abs( hDo * dydx(:) ) + tiny(1.0_D)

                        if( (diffEq%x + hDo - x) * (diffEq%x + hDo - xStart) > 0.0_D ) hDo = x - diffEq%x ! If stepsize overshoots, decrease it:
                        
                        call BurlischStoerStep(diffEq, dydx, yScal, hDo, hDid, hNext, derivatives)
                        
                        if(hDid == hDo) then
                                diffEq%nGood = diffEq%nGood + 1
                        else
                                diffEq%nBad = diffEq%nBad + 1
                        end if

                        if( (diffEq%x - x) * (x - xStart) >= 0.0_D) return ! done

                        if( abs(hNext) < diffEq%hMin ) then
                                write(0, *) 'stepsize smaller than minimum in diffEq integrate.'
                                stop
                        end if

                        hDo = hNext

                end do

                write(0, *) 'WARNING: Too many steps in diffEq integrate.'

        end subroutine integrate

        subroutine BurlischStoerStep(diffEq, dydx, yScal, hTry, hDid, hNext, derivatives)

                use mPolynomialInterpolation
                use mMisc

                type(tDiffEq), intent(inout) :: diffEq
                real(D), dimension( diffEq%nEq ), intent(in) :: dydx, yScal
                real(D), intent(in) :: hTry
                real(D), intent(out) :: hDid, & ! stepsize accomplished
                                        hNext   ! suggested next stepsize
                interface
                        function derivatives(diffEq, x, y)
                                use mPrecision
                                use mObject
                                type(tDiffEq), intent(in) :: diffEq
                                real(D), intent(in) :: x, y(:)
                                real(D) :: derivatives( size(y) )
                        end function derivatives
                end interface
                
!                real(D), dimension( size(diffEq%y) ) :: yExtrapolated
                real(D), dimension( diffEq%nEq ) :: yExtrapolated
!                real(D), dimension( MAX_ROWS, size(diffEq%y) ) :: ySeq
                real(D), dimension( MAX_ROWS, diffEq%nEq ) :: ySeq
                integer :: i, k
!                real(D) :: xNew, h, yErrs( size(diffEq%y) )
                real(D) :: xNew, h, yErrs( diffEq%nEq )
                real(D) :: stepsize2(MAX_ROWS), err(2:MAX_ROWS)
                real(D) :: maxError, stepsizeFactor, scale, workMin
                logical :: didReduction

                h = hTry

                main_loop: do i = 1, MAX_STEPS

                        stepsize_seq: do k = 1, diffEq%k_max

                                xNew = diffEq%x + h

!                                write(0, *) xNew, diffEq%x, h
                                call assert(xNew /= diffEq%x, 'step size underflow in BurlischStoerStep.')

                                call modifiedMidpoint(diffEq, dydx, h, SEQUENCE(k), ySeq(k, :), derivatives)

                                ! We get to square the extrapolation variable
                                ! because the error series is even (see
                                ! Numerical Recipes Eq. (16.3.3)
                                stepsize2(k) = ( h / SEQUENCE(k) ) ** 2

                                call extrapolate( stepsize2(1:k), ySeq(1:k, :), yExtrapolated, yErrs )

                                if(k == 1) cycle

                                maxError = maxval(abs( yErrs(:) / yScal(:) )) ! Scale errors according to yScal:

                                maxError = max( tiny(1.0_D), maxError ) / diffEq%eps ! Scale relative to tolerance

                                err(k) = (maxError * SAFETY_FACTOR1) ** ( 1.0_D / real(2 * k - 1, D) ) ! Compute error estimate using NR Eq. (16.4.17) but for k = ROW NUMBER

                                ! We want to test for convergence in an order window, i.e., NR Eq. (16.4.16), but also in case this is the first attempt:
                                if(k >= diffEq%k_opt - 1 .or. diffEq%first) then

                                        if(maxError < 1.0_D) then ! Converged
                                                exit main_loop
                                        else
                                                if( reduction(diffEq, k, err, stepsizeFactor) ) exit stepsize_seq
                                        end if

                                end if

                        end do stepsize_seq

                        h = h * stepsizeFactor
                        didReduction = .true.

                end do main_loop

                if(i >= MAX_STEPS) write(0, *) 'WARNING: MAX_STEPS exceeded in BurlischStoerStep'

                diffEq%x = xNew ! Successful step taken.
!                diffEq%y = yExtrapolated
                diffEq%y(1:diffEq%nEq) = yExtrapolated
                hDid = h
                if(diffEq%first) diffEq%first = .false.

                diffEq%k_opt = 1 + iminloc(diffEq%A(2:k) * max( err(2:k), SCALE_MAX )) ! Compute optimal row for convergence as in NR Eq. (16.4.9)
                scale = max( err(diffEq%k_opt), SCALE_MAX )
                workMin = scale * diffEq%A(diffEq%k_opt + 1)
                hNext = h / scale
                if(diffEq%k_opt >= k .and. diffEq%k_opt /= diffEq%k_max .and. .not. didReduction) then
                        stepsizeFactor = max( scale / diffEq%alpha(diffEq%k_opt - 1, diffEq%k_opt), SCALE_MAX ) ! Eq. (16.4.11):

                        if( diffEq%A(diffEq%k_opt + 1) * stepsizeFactor <= workMin ) then ! This is the condition of Eq. (16.4.12):
                                hNext = h / stepsizeFactor
                                diffEq%k_opt = diffEq%k_opt + 1
                        end if
                end if

        end subroutine BurlischStoerStep

        subroutine modifiedMidpoint(diffEq, dydx, h, nSteps, yout, derivatives)

                type(tDiffEq), intent(in) :: diffEq
!                real(D), intent(in) :: dydx( size(diffEq%y) ), h
                real(D), intent(in) :: dydx( diffEq%nEq ), h
                integer, intent(in) :: nSteps
                real(D), intent(out) :: yout(:)
                interface
                        function derivatives(diffEq, x, y)
                                use mPrecision
                                use mObject
                                type(tDiffEq), intent(in) :: diffEq
                                real(D), intent(in) :: x, y(:)
                                real(D) :: derivatives( size(y) )
                        end function derivatives
                end interface

                real(D), target, dimension( size(yout) ) :: zScratch1, zScratch2, zScratch3
                real(D), pointer, dimension(:) :: z_mMinus1, z_m, z_mPlus1
                real(D) :: subH
                integer :: m

                nullify(z_mMinus1, z_m, z_mPlus1)

                z_mMinus1 => zScratch1
                z_m => zScratch2
                z_mPlus1 => zScratch3

                subH = h / real(nSteps, D)

                ! Initial values (m = 1):
!                z_mMinus1 = diffEq%y
                z_mMinus1 = diffEq%y(1:diffEq%nEq)
                z_m = z_mMinus1 + subH * dydx

                do m = 1, nSteps - 1

                        z_mPlus1 = z_mMinus1 + 2.0_D * subH * derivatives( diffEq, diffEq%x + real(m, D) * subH, z_m )

                        call permute(z_mMinus1, z_m, z_mPlus1)

                end do
                ! now m = n

                yout = 0.5_D * ( z_m + z_mMinus1 + subH * derivatives(diffEq, diffEq%x + h, z_m) )

        end subroutine modifiedMidpoint

        ! Check for stepsize reduction. The conditions and reduction factors
        ! are taken from Numerical Recipes in Fortran 90, chapter B16 (p. 1304):
        logical function reduction(diffEq, k, err, reductionFactor)
                type(tDiffEq), intent(in) :: diffEq
                integer, intent(in) :: k
                real(D), intent(in) :: err(2:)
                real(D), intent(out) :: reductionFactor

                if(k == diffEq%k_max .or. k == diffEq%k_opt + 1) then

                        reductionFactor = SAFETY_FACTOR2 / err(k)
                        reduction = .true.

                else if(k == diffEq%k_opt) then

                        if( diffEq%alpha(diffEq%k_opt - 1, diffEq%k_opt) < err(k) ) then
                                reductionFactor = 1.0_D / err(k)
                                reduction = .true.
                        else
                                reduction = .false.
                        end if

                else if(diffEq%k_opt == diffEq%k_max) then

                        if( diffEq%alpha(k - 1, diffEq%k_max - 1) < err(k) ) then
                                reductionFactor = diffEq%alpha(k - 1, diffEq%k_max - 1) * SAFETY_FACTOR2 / err(k)
                                reduction = .true.
                        else
                                reduction = .false.
                        end if

                else if( diffEq%alpha(k - 1, diffEq%k_opt) < err(k) ) then

                        reductionFactor = diffEq%alpha(k - 1, diffEq%k_opt - 1) / err(k)
                        reduction = .true.
                else
                        reduction = .false.
                end if

                if(reduction) reductionFactor = max( min(reductionFactor, REDUCTION_MIN), REDUCTION_MAX )

        end function reduction

!        type(tUniverse) function getUniverse_diffEq(diffEq)
!
!                type(tDiffEq), intent(in) :: diffEq
!
!                getUniverse_diffEq = diffEq%u
!
!        end function getUniverse_diffEq

        subroutine permute(p1, p2, p3)

                real(D), pointer, dimension(:) :: p1, p2, p3
                real(D), pointer, dimension(:) :: temp

                temp => p1
                p1 => p2
                p2 => p3
                p3 => temp

        end subroutine permute

end module mDiffEq
